import os
import pandas as pd
import adata
from utils.logger import logger

class StockDataFetcher:
    """
    用于获取和存储股票数据的类。
    
    该类允许从特定数据源获取股票数据，并将数据存储为pickle文件以供后续使用。
    """

    def __init__(self, stock_code,start_date='2024-10-01',forced_update=False, data_dir='./data'):
        """
        初始化StockDataFetcher对象。
        
        参数:
        stock_code: str - 股票代码。
        data_dir: str - 存储数据的目录，默认为当前目录下的'data'目录。
        """
        self.stock_code = stock_code
        self.start_date=start_date
        self.forced_update = forced_update
        self.data_dir = data_dir
        # 拼接pickle文件的完整路径
        self.pickle_file_path = os.path.join(self.data_dir, f'{self.stock_code}.pkl')

    def safe_read_pickle(self):
        """
        安全地读取pickle文件。
        
        尝试读取pickle文件，如果失败则记录错误并返回None。
        
        返回:
        pd.DataFrame or None - 读取成功的DataFrame对象或None。
        """
        """安全地读取pickle文件"""
        try:
            return pd.read_pickle(self.pickle_file_path)
        except Exception as e:
            logger.error(f"读取pickle文件时发生错误: {e}")
            return None

    def safe_write_pickle(self, dataframe):
        """
        安全地写入pickle文件。
        
        尝试将DataFrame写入pickle文件，如果失败则记录错误。
        
        参数:
        dataframe: pd.DataFrame - 要写入的DataFrame对象。
        """
        """安全地写入pickle文件"""
        try:
            dataframe.to_pickle(self.pickle_file_path)
        except Exception as e:
            logger.error(f"保存pickle文件时发生错误: {e}")

    def validate_dataframe(self, df):
        """
        验证DataFrame的有效性。
        
        检查DataFrame是否为空，以及是否包含预期的列。
        
        参数:
        df: pd.DataFrame - 要验证的DataFrame对象。
        
        返回:
        bool - 如果DataFrame有效则为True，否则为False。
        """
        """验证DataFrame是否符合预期"""
        if df is None or df.empty:
            logger.warning("数据为空或未定义")
            return False
        # 假设我们期望DataFrame有某些特定的列，这里进行检查
        expected_columns = ['close', 'open']  
        for column in expected_columns:
            if column not in df.columns:
                logger.error(f"缺少预期的列: {column}")
                return False
        return True

    def fetch_stock_data(self, k_type=1):
        """
        获取股票数据。
        
        从数据源获取指定股票的市场数据。
        
        参数:
        start_date: str - 开始日期，格式为'YYYY-MM-DD'。
        k_type: int - K线类型，具体取决于数据源的定义。
        
        返回:
        pd.DataFrame or None - 获取成功的DataFrame对象或None。
        """
        """获取股票数据（前复权）"""
        try:
            df = adata.stock.market.get_market(stock_code=self.stock_code, k_type=k_type, start_date=self.start_date)
            if self.validate_dataframe(df):
                return df
            else:
                return None
        except Exception as e:
            logger.error(f"获取股票数据时发生错误: {e}")
            return None

    def fetch_or_load_data(self):
        """
        如果存在，则加载pickle文件中的数据；否则尝试获取数据，并将其保存为pickle文件。
        
        返回:
        pd.DataFrame or None - 加载或获取的数据的DataFrame对象，如果无法获取或验证数据则为None。
        """
        """检查pickle文件是否存在，如果存在则读取，否则尝试获取数据并保存"""
        if self.forced_update:
             logger.info("强制更新数据")
             os.remove(self.pickle_file_path)
             
        if os.path.exists(self.pickle_file_path):
            df = self.safe_read_pickle()
        else:
            df = self.fetch_stock_data()
            if df is not None:
                self.safe_write_pickle(df)
            else:
                logger.error("无法获取或验证数据")
                return None
        return df